#pragma warning(disable : 4355)
#include "Avalon/Threading/ThreadPool.hpp"
#include <boost/date_time/posix_time/posix_time.hpp>

using namespace Avalon;
using namespace Avalon::Threading;
using namespace boost;
using namespace boost::posix_time;
using namespace std;

namespace {
  const time_duration LOWER_BOUND_WAIT_TIME = seconds(30);
  const time_duration UPPER_BOUND_WAIT_TIME = seconds(60);
}

ThreadPool::BaseTask::BaseTask() {}

ThreadPool::BaseTask::~BaseTask() {}

class ThreadPool::TaskThread {
  public:
    TaskThread(ThreadPool& threadPool);

    ~TaskThread();

    bool SetTask(BaseTask* task);

    void Run();

    void Stop();

  private:
    mutex m_mutex;
    ThreadPool& m_threadPool;
    BaseTask* m_task;
    bool m_available;
    bool m_stopped;
    scoped_ptr<thread> m_thread;
    function<void ()> m_runFunction;
    condition_variable m_taskAvailableCondition;
};

ThreadPool::TaskThread::TaskThread(ThreadPool& threadPool)
    : m_threadPool(threadPool),
      m_task(NULL),
      m_available(true),
      m_stopped(false),
      m_runFunction(bind(&TaskThread::Run, this)) {}

ThreadPool::TaskThread::~TaskThread() {}

bool ThreadPool::TaskThread::SetTask(BaseTask* task) {
  lock_guard<mutex> lock(m_mutex);
  assert(m_task == NULL);
  if(!m_available) {
    return false;
  }
  m_task = task;
  if(m_thread == NULL) {
    m_thread.reset(new thread(m_runFunction));
  } else {
    m_taskAvailableCondition.notify_one();
  }
  return true;
}

void ThreadPool::TaskThread::Run() {
  while(true) {
    {
      unique_lock<mutex> lock(m_mutex);
      if(m_task == NULL) {
        time_duration waitTime = LOWER_BOUND_WAIT_TIME + seconds(rand() %
          (UPPER_BOUND_WAIT_TIME - LOWER_BOUND_WAIT_TIME).total_seconds());
        if(!m_available ||
            !m_taskAvailableCondition.timed_wait(lock, waitTime) ||
            m_task == NULL) {
          m_available = false;
          bool stopped = m_stopped;
          lock.unlock();
          if(!stopped) {
            m_threadPool.RemoveThread(this);
          }
          delete this;
          return;
        }
      }
      m_task->Run();
      delete m_task;
      m_task = NULL;
    }
    if(!m_threadPool.AddThread(this)) {
      delete this;
      return;
    }
  }
}

void ThreadPool::TaskThread::Stop() {
  lock_guard<mutex> lock(m_mutex);
  m_available = false;
  m_stopped = true;
  m_task = NULL;
  m_taskAvailableCondition.notify_one();
}

ThreadPool::ThreadPool(int maxThreadCount)
    : m_maxThreadCount(maxThreadCount),
      m_runningThreads(0) {}

ThreadPool::~ThreadPool() {
  unique_lock<mutex> lock(m_mutex);
  m_maxThreadCount = 0;
  while(!m_threads.empty()) {
    TaskThread* taskThread = m_threads.back();
    taskThread->Stop();
    m_threads.pop_back();
  }
  if(m_runningThreads != 0) {
    m_threadsFinishedCondition.wait(lock);
  }
}

int ThreadPool::GetMaxThreadCount() const {
  return m_maxThreadCount;
}

void ThreadPool::QueueTask(BaseTask* task) {
  lock_guard<mutex> lock(m_mutex);
  bool threadFound = false;
  while(!threadFound) {
    if(m_threads.empty() && m_runningThreads < m_maxThreadCount) {
      TaskThread* taskThread = new TaskThread(*this);
      taskThread->SetTask(task);
      ++m_runningThreads;
      return;
    } else if(!m_tasks.empty() || m_threads.empty()) {
      m_tasks.push_back(task);
      return;
    }
    TaskThread* taskThread = m_threads.front();
    threadFound = taskThread->SetTask(task);
    if(threadFound) {
      ++m_runningThreads;
      m_threads.pop_front();
    }
  }
}

bool ThreadPool::AddThread(TaskThread* taskThread) {
  lock_guard<mutex> lock(m_mutex);
  if(m_maxThreadCount == 0) {
    --m_runningThreads;
    if(m_runningThreads == 0) {
      m_threadsFinishedCondition.notify_one();
    }
    return false;
  }
  if(m_tasks.empty()) {
    --m_runningThreads;
    m_threads.push_back(taskThread);
    return true;
  }
  BaseTask* task = m_tasks.front();
  bool taskIsAssigned = taskThread->SetTask(task);
  if(taskIsAssigned) {
    m_tasks.pop_front();
  }
  return true;
}

void ThreadPool::RemoveThread(TaskThread* taskThread) {
  lock_guard<mutex> lock(m_mutex);
  --m_runningThreads;
  if(m_maxThreadCount == 0 && m_runningThreads == 0) {
    m_threadsFinishedCondition.notify_one();
  }
}
